import os
import glob
import random, csv

def load_csv(root, filename, name2label):
    '''
    从csv文件返回 images, labels列表
    :param root: 数据集根目录
    :param filename: csv文件名
    :param name2label: 类别名编码表
    :return:返回图片路径list和标签list
    '''
    if not os.path.exists(os.path.join(root, filename)):
        # 如果csv文件不存在，则创建
        images = []
        for name in name2label.keys():# 遍历所有子目录，获得所有图片
            # 只考虑后缀为png, jpg, jpeg的图片：’pokemon\\mewtwo\\00001.png‘
            images += glob.glob(os.path.join(root, name, '*.png'))
            images += glob.glob(os.path.join(root, name, '*.jpg'))
            images += glob.glob(os.path.join(root, name, '*.jpeg'))
        random.shuffle(images) # 随机打散
        # 创建csv文件， 并存储图片路径及其label信息
        with open(os.path.join(root, filename), mode='w', newline='') as f:
            writer = csv.writer(f)
            for img in images: # 'pokemon\\bulbasaur\\00000000.png'
                name = img.split(os.sep)[-2]
                # print(name)
                label = name2label[name]
                # 'pokemon\\bulbasaur\\00000000.png'
                writer.writerow([img, label])
            print('written into cscl file:', filename)
    images, labels = [], []
    with open(os.path.join(root, filename)) as f:
        reader = csv.reader(f)
        for row in reader:
            # 'pokemon\\bulbasaur\\00000000.png', 0
            img, label = row
            label = int(label)
            images.append(img)
            labels.append(label)
    # 返回图片路径list和标签list
    return images, labels

def load_pokemon(root, mode= 'train'):
    '''
    加载宝可梦数据集，并做划分，前60%为训练集， 60%-80%为验证集， 80%-100%为测试集
    :param root: 数据集根目录
    :param mode: 数据集的应用模式
    :return: 返回对应模式的数据和标签
    '''
    # 创建数字编码表（字典）
    name2label = {}
    # 遍历根目录下所有的子文件夹，排序，保证映射关机
    for name in sorted(os.listdir(os.path.join(root))):
        # 跳过非文件夹对象
        if not os.path.isdir(os.path.join(root, name)):
            continue
        # 给每个类别编码一个数字
        name2label[name] = len(name2label.keys())
    images, labels = load_csv(root, 'images.csv', name2label)
    # 数据集划分
    if mode == 'train': # 60%
        images = images[:int(0.6 * len(images))]
        labels = labels[:int(0.6 * len(labels))]
    elif mode == 'val': # 20%
        images = images[int(0.6 * len(images)):int(0.8 * len(images))]
        labels = labels[int(0.6 * len(labels)):int(0.8 * len(labels))]
    else: # 20%
        images = images[int(0.8 * len(images)):]
        labels = labels[int(0.8 * len(labels)):]

    return images, labels, name2label

def main():
    images, labels, name2label = load_pokemon(r'D:\下载（文档）\pokeman')
    print(images)
    print(labels)
    print(name2label)

if __name__ == '__main__':
    main()









